import pickle
import  torch
import pandas as pd
from scipy import sparse
import dgl
import pkuseg
import  numpy as np

class RumorGraph(object):
    def __init__(self,
                 train_file="./data/train_1_3_sim093.csv",
                 dev_file="./data/dev_1_3_sim093.csv",
                 test_file="./data/test_dataset.csv",
                 test_semant="./data/test_semant.pkl",
                 dev_semant="./data/dev_semant_1_3.pkl",
                 train_semant="./data/train_semant_1_3.pkl",
                 edge_sim = 0.98
                 ):
        self.g = dgl.DGLGraph()
        self.dev_file = dev_file
        self.train_file = train_file
        self.test_file = test_file
        self.test_semant = test_semant
        self.dev_semant = dev_semant
        self.train_semant = train_semant
        self.edge_sim = edge_sim
        self.seg = pkuseg.pkuseg()
        with open("./model/emb_layer/word_dic.pkl", "rb") as fr:
            self.dic = pickle.load(fr)
        self.train_num = 0
        self.vlaid_num = 0
        self.test_num = 0

    def ObtainScoMtx(self):
        with open(self.train_semant, "rb") as fr:
            tr = pickle.load(fr)
            self.train_num = tr.size(0)
        with open(self.dev_semant, "rb") as fr:
            dev = pickle.load(fr)
            self.valid_num = dev.size(0)
        with open(self.test_semant, "rb") as fr:
            te = pickle.load(fr)
            self.test_num = te.size(0)
        vecs = torch.cat([tr, dev, te], dim=0)
        dot = torch.matmul(vecs, vecs.transpose(0, 1))
        norm_vecs = vecs.norm(p=2, dim=1).unsqueeze(0)
        norm_mtx = torch.matmul(norm_vecs.transpose(0, 1), norm_vecs)
        sco = dot / norm_mtx
        return sco, vecs

    def obtain_labels(self, csv_file):
        df = pd.read_csv(csv_file)
        return torch.tensor(df[['ncw_label', 'fake_label', 'real_label']].values.argmax(axis=1))

    def init_node_feats(self, feats):
        self.g.ndata['feat'] = feats

    def init_edges(self, sim_mtx):
        sco = sim_mtx.ge(self.edge_sim).int().cpu()
        sr = sparse.coo_matrix((sco - torch.eye(sco.size(0))).numpy())
        src = sr.row
        dst = sr.col
        self.g.add_edges(src, dst)
        self.g.add_edges(dst, src)

    def init_labels(self):
        size_0 = self.g.ndata['feat'].size(0)
        l = torch.cat([self.obtain_labels(self.train_file), self.obtain_labels(self.dev_file)])
        size_re = size_0 - l.size(0)
        assert size_re == self.test_num
        self.g.ndata['label'] = torch.cat([l, torch.ones(size_re, dtype=torch.int64 ) *-1], dim=0)

    def obtain_word_ids(self, csv_file):
        contents = pd.read_csv(csv_file)['content'].values.tolist()
        word_ids = np.array([np.array([self.dic[w] for w in sent if w in self.dic]) for sent in contents])
        return word_ids

    def init_node_contents(self):
        self.g.ndata['words'] = np.concatenate([self.obtain_word_ids(self.train_file), self.obtain_word_ids(self.dev_file), self.obtain_word_ids(self.test_file)])

    def graph_construtor(self):
        sco, vecs = self.ObtainScoMtx()
        self.g.add_nodes(vecs.size(0))
        self.init_node_feats(vecs)
        del vecs
        self.init_edges(sco)
        del sco
        self.init_labels()
        self.init_node_contents()
        print('the nodes number is:', self.g.number_of_nodes())
        print('the edge numbers:', self.g.number_of_edges())
        print('the train num:', self.train_num)
        print('the valid num:', self.valid_num)
        print('the test num:', self.test_num)